package org.glandais.profiler;

import java.io.IOException;
import java.util.HashSet;
import java.util.Set;

import javassist.CannotCompileException;
import javassist.ClassPath;
import javassist.ClassPool;
import javassist.CtBehavior;
import javassist.CtClass;
import javassist.CtMethod;
import javassist.LoaderClassPath;
import javassist.Modifier;
import javassist.NotFoundException;
import javassist.bytecode.MethodInfo;

public class ClassFileTransformer implements
		java.lang.instrument.ClassFileTransformer {

	private ClassPool classPool = ClassPool.getDefault();
	private Set<ClassLoader> loaders = new HashSet<ClassLoader>();
	private CtClass ctThrowable = classPool.get("java.lang.Throwable");

	public static boolean doTransformClasses = true;

	class BehaviorInfo {

		String clazz = null;
		String methodSignature = null;
		String sig = null;
		String sthis = "null";
		String type = "UNKN";

		public BehaviorInfo(CtBehavior method) {
			super();
			MethodInfo methodInfo2 = method.getMethodInfo2();
			clazz = method.getDeclaringClass().getName();
			methodSignature = method.getSignature();
			if (methodInfo2.isConstructor()) {
				type = "CONS";
				sig = "new " + clazz + " " + methodSignature;
			} else if (methodInfo2.isStaticInitializer()) {
				type = "INIT";
				sig = clazz + "{} " + methodSignature;
			} else if (Modifier.isStatic(method.getModifiers())) {
				type = "STAT";
				sig = "static " + clazz + "." + method.getName() + " "
						+ methodSignature;
			} else if (method instanceof CtMethod) {
				type = "METH";
				sthis = "$0";
				sig = clazz + "." + method.getName() + " " + methodSignature;
			}
		}

	}

	public ClassFileTransformer(String agentArgument) throws Exception {
		super();
	}

	public final byte[] loadAndTransform(ClassLoader loader, String className,
			Class<?> clazz, java.security.ProtectionDomain domain)
			throws IOException, CannotCompileException, NotFoundException {
		if (!loaders.contains(loader)) {
			addClassLoader(loader);
		}
		CtClass ctClass = classPool.get(className);
		return doClass(loader, className, clazz, ctClass);
	}

	public final boolean shouldTransform(String className) {
		if (doTransformClasses) {
			for (String excluded : Agent.configBean.getExcluded()) {
				if (className.replace('/', '.').startsWith(
						excluded.replace('/', '.'))) {
					if (Agent.configBean.isLogTraceInstrumentation()) {
						Agent.configBean
								.logTraceInstrumentation("-- Excluded : "
										+ className);
					}
					return false;
				}
			}
			if (className.replace('/', '.').startsWith("org.glandais.profiler")) {
				if (Agent.configBean.isLogTraceInstrumentation()) {
					Agent.configBean.logTraceInstrumentation("-- Excluded : "
							+ className);
				}
				return false;
			}
			if (Agent.configBean.getIncluded().size() == 0) {
				return true;
			} else {
				for (String included : Agent.configBean.getIncluded()) {
					if (className.replace('/', '.').startsWith(
							included.replace('/', '.'))) {
						return true;
					}
				}
			}
		}
		return false;
	}

	@Override
	public final byte[] transform(ClassLoader loader, String className,
			Class<?> clazz, java.security.ProtectionDomain domain, byte[] bytes) {
		TraceRecorder.totalTimeTL.get().startProcessing();
		if (shouldTransform(className)) {
			byte[] result = doClass(loader, className, clazz, bytes);
			TraceRecorder.totalTimeTL.get().stopProcessing();
			return result;
		}
		TraceRecorder.totalTimeTL.get().stopProcessing();
		return bytes;
	}

	private synchronized final void addClassLoader(ClassLoader loader) {
		if (!loaders.contains(loader)) {
			if (Agent.configBean.isLogTraceInstrumentation()) {
				Agent.configBean
						.logTraceInstrumentation("++ Adding class loader : "
								+ loader);
			}
			loaders.add(loader);
			ClassPath cp = new LoaderClassPath(loader);
			classPool.appendClassPath(cp);
		}
	}

	private final byte[] doClass(ClassLoader loader, String name,
			Class<?> clazz, byte[] b) {
		CtClass cl = null;
		try {
			cl = classPool.makeClass(new java.io.ByteArrayInputStream(b));
			return doClass(loader, name, clazz, cl);
		} catch (Throwable e) {
			Agent.configBean.logError("EE Could not instrument  " + name, e);
		}
		return new byte[0];
	}

	private final byte[] doClass(ClassLoader loader, String name,
			Class<?> clazz, CtClass cl) {
		if (!loaders.contains(loader)) {
			addClassLoader(loader);
		}

		if (Agent.configBean.isLogVerbose()) {
			Agent.configBean.logVerbose("++ Transforming : " + name);
		}

		byte[] b = new byte[0];
		try {
			if (cl.isInterface() == false) {
				CtBehavior[] methods = cl.getDeclaredBehaviors();
				for (int i = 0; i < methods.length; i++) {
					if (methods[i].isEmpty() == false) {
						BehaviorInfo behaviorInfo = new BehaviorInfo(methods[i]);
						doMethod(behaviorInfo, methods[i]);
					}
				}
				b = cl.toBytecode();
			}
		} catch (Throwable e) {
			Agent.configBean.logError("EE Could not instrument  " + name, e);
		} finally {
			if (cl != null) {
				cl.detach();
			}
		}
		return b;
	}

	private final void doMethod(BehaviorInfo behaviorInfo, CtBehavior method)
			throws NotFoundException, CannotCompileException {
		method.insertBefore("org.glandais.profiler.TraceRecorder.enter( \""
				+ behaviorInfo.type + "\", " + behaviorInfo.sthis + " ,\""
				+ behaviorInfo.sig + "\" );");
		String exitCode = "org.glandais.profiler.TraceRecorder.exit( \""
				+ behaviorInfo.type + "\", " + behaviorInfo.sthis + " , \""
				+ behaviorInfo.sig + "\" );";
		method.addCatch("{ " + exitCode + " throw $e; }", ctThrowable);
		method.insertAfter(exitCode);
	}
}
